import tensorflow as tf
from model.yoloV3 import YOLOv3,OutputParser
from utils.config_utils import parse_anchors
import numpy as np

# 损失函数定义
def Loss(img_shape, class_num=80):
    # 设置anchor
    YOLOv3_anchor = parse_anchors("config/anchor")
    anchors = dict(zip(reversed(range(3)),
                        [YOLOv3_anchor[0:3].tolist(), YOLOv3_anchor[3:6].tolist(), YOLOv3_anchor[6:9].tolist()]))
    # 设置Loss的输入
    input_shapes = [
        (img_shape[0] // 32, img_shape[1] // 32, 3, 5 + class_num),
        (img_shape[0] // 16, img_shape[1] // 16, 3, 5 + class_num),
        (img_shape[0] // 8, img_shape[1] // 8, 3, 5 + class_num)
    ]
    # 定义Loss的inputs:预测值和目标值
    inputs = [tf.keras.Input(input_shape) for input_shape in input_shapes]
    labels = [tf.keras.Input(input_shape) for input_shape in input_shapes]
    # 计算损失
    losses = list()
    # 遍历多个尺度
    for l in range(3):
        input_shape_of_this_layer = input_shapes[l]
        anchors_of_this_layer = anchors[l]
        input_of_this_layer = inputs[l]
        label_of_this_layer = labels[l]
        # 网络输出
        pred_xy, pred_wh, pred_box_confidence, pred_class = OutputParser(
            input_shape_of_this_layer, img_shape, anchors_of_this_layer)(input_of_this_layer)
        # 预测框
        pred_box = tf.keras.layers.Concatenate()([pred_xy, pred_wh])
        # 真实值
        true_box = tf.keras.layers.Lambda(
            lambda x: x[..., 0:4])(label_of_this_layer)
        true_box_confidence = tf.keras.layers.Lambda(
            lambda x: x[..., 4])(label_of_this_layer)
        true_class = tf.keras.layers.Lambda(
            lambda x: x[..., 5:])(label_of_this_layer)
        # 正负样本
        object_box = tf.keras.layers.Lambda(
            lambda x: tf.cast(x, dtype=tf.bool))(true_box_confidence)
        # 1.box:MSE,只有正样本参与损失计算
        pos_loss = tf.keras.layers.Lambda(lambda x: tf.math.reduce_sum(tf.keras.losses.MSE(
            tf.boolean_mask(x[0], x[2]), tf.boolean_mask(x[1], x[2]))))([true_box, pred_box, object_box])
        # 2.confidence:正负样本都要计算，二分类的交叉熵损失
        confidence_loss = tf.keras.layers.Lambda(lambda x: tf.keras.losses.BinaryCrossentropy(from_logits=False)(tf.boolean_mask(x[0], x[2]), tf.boolean_mask(x[1], x[2])) + 0.5 * tf.keras.losses.BinaryCrossentropy(
            from_logits=False)(tf.boolean_mask(x[0], tf.math.logical_not(x[2])), tf.boolean_mask(x[1], tf.math.logical_not(x[2]))))([true_box_confidence, pred_box_confidence, object_box])
        # 3.class:只有正样本计算损失 二分类的交叉熵损失
        class_loss = tf.keras.layers.Lambda(lambda x: tf.keras.losses.BinaryCrossentropy(from_logits=False)(
            tf.boolean_mask(x[0], x[2]), tf.boolean_mask(x[1], x[2])))([true_class, pred_class, object_box])
        # 4.总损失
        loss = tf.keras.layers.Lambda(lambda x: tf.math.add_n(x))(
            [pos_loss, confidence_loss, class_loss])
        losses.append(loss)
    loss = tf.keras.layers.Lambda(lambda x: tf.math.add_n(x))(losses)
    return tf.keras.Model(inputs=(*inputs, *labels), outputs=loss)

if __name__=="__main__":
    # 设置anchor
    anchors = {2: [[5, 5], [6, 7], [7, 9]], 1: [[10, 11], [
        13, 15], [19, 21]], 0: [[27, 31], [43, 50], [79, 93]]}
    print(anchors)
    YOLOv3_anchor = parse_anchors("../config/anchor")
    print(YOLOv3_anchor[0:3])
    anchors2 = dict(zip(reversed(range(3)),[YOLOv3_anchor[0:3].tolist(),YOLOv3_anchor[3:6].tolist(),YOLOv3_anchor[6:9].tolist()]))
    print(anchors2)


